import torch
import torch.nn as nn
import torch.nn.functional as F
from src.Main.config import SOS_token, device, attn_weights_list


class Attn(nn.Module):
    def __init__(self, method, hidden_size):
        """

        :param method: 'dot',采用dot_score函数计算隐状态和encoder输出相似度得分
        :param hidden_size: 500 隐藏层的大小
        """
        super(Attn, self).__init__()
        self.method = method
        if self.method not in ['dot', 'general', 'concat']:
            raise ValueError(self.method, "is not an appropriate attention method.")
        self.hidden_size = hidden_size
        if self.method == 'general':
            self.attn = nn.Linear(self.hidden_size, hidden_size)
        elif self.method == 'concat':
            self.attn = nn.Linear(self.hidden_size * 2, hidden_size)
            # 缓存临时状态，不让临时变量注册为模型变量
            self.v = nn.parameter(torch.FloatTensor(hidden_size))

    def dot_score(self, hidden, encoder_output):
        """
        :param hidden:Tensor:(1,256,500)
        :param encoder_output:Tensor:(10,256,500)
        :return:
        """
        return torch.sum(hidden * encoder_output, dim=2)

    def general_score(self, hidden, encoder_output):
        """
        :param hidden:Tensor:(1,256,500)
        :param encoder_output:Tensor:(10,256,500)
        :return:
        """
        energy = self.attn(encoder_output)
        return torch.sum(hidden * energy, dim=2)

    def concat_score(self, hidden, encode_output):
        """

        :param hidden:Tensor:(1,256,500)
        :param encode_output:Tensor:(10,256,500)
        :return:
        """
        energy = self.attn(torch.cat((hidden.expand(encode_output.size(0), -1, -1), encode_output), 2)).tanh()
        return torch.sum(self.v * energy, dim=2)

    def forward(self, hidden, encoder_outputs):
        # 根据所指定的方法计算隐状态和encoder输出的相似度的得分
        if self.method == 'general':
            attn_energies = self.general_score(hidden, encoder_outputs)
        elif self.method == 'concat':
            attn_energies = self.concat_score(hidden, encoder_outputs)
        elif self.method == 'dot':
            attn_energies = self.dot_score(hidden, encoder_outputs)

        # 将得到的矩阵进行转置
        attn_energies = attn_energies.t()

        # 将score得分归一化后计算得到的概率增加一个维度并返回
        return F.softmax(attn_energies, dim=1).unsqueeze(1)


class LuongAttnDecoderRNN(nn.Module):
    def __init__(self, attn_model, embedding, hidden_size, output_size, n_layers=1, drop_out=0.1):
        """

        :param attn_model:'dot' #选用dot_score方法计算概率分数
        :param embedding:（7826,500）
        :param hidden_size: 500 #隐藏层一层的神经元个数
        :param output_size:7826
        :param n_layers:2
        :param drop_out:0.1 #防止模型过拟合
        """
        super(LuongAttnDecoderRNN, self).__init__()

        self.attn_model = attn_model
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.n_layers = n_layers
        self.dropout = drop_out

        self.embedding = embedding
        # drop_out=0.1,每个神经元有0.1的可能性不被激活，该函数作用是防止过拟合
        self.embedding_dropout = nn.Dropout(drop_out)
        self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=(0 if n_layers == 1 else drop_out))
        self.concat = nn.Linear(hidden_size * 2, hidden_size)
        self.out = nn.Linear(hidden_size, output_size)

        # 把attn前向传播的结果进行接收
        self.attn = Attn(attn_model, hidden_size)

    def forward(self, input_step, last_hidden, encoder_outputs):
        """

        :param input_step:Tensor(1,256)
        :param last_hidden:Tensor(2,256,500)
        :param encoder_outputs:Tensor(10,256,500)
        :return:
        """
        embedded = self.embedding(input_step)
        # 添加dropout防止过拟合
        embedded = self.embedding_dropout(embedded)
        # 将数据传入gru得到decoder模型的输出和decoder模型的隐状态
        rnn_output, hidden = self.gru(embedded, last_hidden)
        # attn_weight:Tensor（256,1，10），将各个输出赋予不一样的权重值
        attn_weights = self.attn(rnn_output, encoder_outputs)
        attn_weights_list.append(attn_weights)
        # 矩阵转置
        context = attn_weights.bmm(encoder_outputs.transpose(0, 1))
        # 压缩一个维度
        rnn_output = rnn_output.squeeze(0)  # Tensor（1,256,500）
        context = context.squeeze(1)
        # 拼接张量，rnn输出和上下文关系进行拼接得到concat_input
        concat_input = torch.cat((rnn_output, context), 1)
        # concat_input进行Linear全连接后进行tanh函数映射
        concat_output = torch.tanh(self.concat(concat_input))
        # 再次进行一个全连接
        output = self.out(concat_output)
        # softmax概率映射
        output = F.softmax(output, dim=1)

        return output, hidden


# 贪婪搜索解码器
class GreedySearchDecoder(nn.Module):
    def __init__(self, encoder, decoder):
        super(GreedySearchDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.att = []

    def forward(self, input_seq, input_length, max_length):
        encoder_outputs, encoder_hidden = self.encoder(input_seq, input_length)
        # encoder编码器的最终隐藏层作为decoder解码器的第一个隐藏输入
        decoder_hidden = encoder_hidden[:self.decoder.n_layers]
        # SOS_token初始化解码器输入
        decoder_input = torch.ones(1, 1, device=device, dtype=torch.long) * SOS_token
        # 初始化张量
        all_tokens = torch.zeros([0], device=device, dtype=torch.long)
        all_scores = torch.zeros([0], device=device)
        # 一次迭代解码一个词标记
        # attn_weights_list.clear()
        for _ in range(max_length):
            decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs)
            # 获取最可能的词标记及其softmax概率分数
            decoder_scores, decoder_input = torch.max(decoder_output, dim=1)
            # 记录token和score
            all_tokens = torch.cat((all_tokens, decoder_input), dim=0)
            all_scores = torch.cat((all_scores, decoder_scores), dim=0)
            # 准备当前token作为下一个解码器输入（需要添加维度）
            decoder_input = torch.unsqueeze(decoder_input, 0)
        self.att = attn_weights_list[:]
        attn_weights_list.clear()
        return all_tokens, all_scores
